"""
| Tools for language models.
"""


import numpy as np


def apply_bert_mask(inputs, pad_mask, tokenizer):
    """
    Apply BERT mask to the token_ids.
    Args:
        token_ids: The list of token ids.
    Returns:
        masked_token_ids: The list of masked token ids.
        labels: The labels for traininig BERT.
    """
    vocab_size = len(tokenizer.vocab)
    bert_mask = np.random.uniform(size=inputs.shape) < 0.15
    bert_mask &= pad_mask

    masked_inputs = inputs * ~bert_mask
    random_uniform = np.random.uniform(size=inputs.shape)
    token_bert_mask = random_uniform < 0.8
    random_bert_mask = random_uniform > 0.9
    true_bert_mask = ~token_bert_mask & ~random_bert_mask

    token_bert_mask = token_bert_mask & bert_mask
    random_bert_mask = random_bert_mask & bert_mask
    true_bert_mask = true_bert_mask & bert_mask

    masked_inputs += tokenizer.mask_token_id * token_bert_mask

    masked_inputs += np.random.randint(0, vocab_size, size=(inputs.shape)) * random_bert_mask
    masked_inputs += inputs * true_bert_mask

    labels = np.where(bert_mask, inputs, -1)

    return masked_inputs, labels